import os

import chromadb
from langchain_community.embeddings import OpenAIEmbeddings


# 创建链接客户端
def create_client():
    # 创建client
    client = chromadb.HttpClient(
        host="127.0.0.1",  # 你的服务器ip
        port=8800,  # 你的服务器端口
    )
    return client


# 创建集合
def create_collection(collection_name: str):
    client = create_client()
    collection = client.get_or_create_collection(
        name=collection_name,  # collection名称
    )
    return collection


# 获取集合
def get_collection(collection_name: str):
    client = create_client()
    collection = client.get_collection(name=collection_name)
    return collection


# 向集合中添加文档
def add_docs(collection_name: str, docs: dict):
    client = create_client()
    collection = client.get_collection(name=collection_name)
    # embedding函数，这里用了langchain的openai embedding函数
    embeddings = OpenAIEmbeddings(openai_api_base='https://api.chatanywhere.tech/v1').embed_documents
    # 向collection中添加文档
    collection.add(
        embeddings=embeddings(texts=["hello world", "goodbye world"]),  # 每个文档的向量
        documents=["hello world", "goodbye world"],  # 文档的原文
        ids=[f"id{i}" for i in range(len(["hello world", "goodbye world"]))]  # 每个文档的 id
    )


# 从数据库中查询数据
def get_data_from_docs(collection_name: str):
    client = create_client()
    collection = client.get_collection(name=collection_name)
    embeddings = OpenAIEmbeddings(openai_api_base='https://api.chatanywhere.tech/v1').embed_documents
    # 查询向量数据库
    result = collection.query(
        n_results=1,
        query_embeddings=embeddings(texts=["hello"]),
    )
    print(result)


def get_opeai_key():
    openai_api_key = os.environ.get('OPENAI_API_KEY')
    print(openai_api_key)
    return ""
